package com.ajocer.springbootinit.aop;

import com.ajocer.springbootinit.annotation.RateLimit;
import com.ajocer.springbootinit.common.ErrorCode;
import com.ajocer.springbootinit.exception.BusinessException;
import io.github.resilience4j.ratelimiter.RateLimiter;
import io.github.resilience4j.ratelimiter.RateLimiterRegistry;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.time.Duration;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 限流切面
 * 基于Resilience4j实现API限流
 */
@Aspect
@Component
@Slf4j
public class CustomRateLimiterAspect {

    @Autowired
    private RateLimiterRegistry rateLimiterRegistry;

    // 缓存已创建的限流器
    private final ConcurrentHashMap<String, RateLimiter> rateLimiterCache = new ConcurrentHashMap<>();

    /**
     * 环绕通知处理限流逻辑
     */
    @Around("@annotation(com.ajocer.springbootinit.annotation.RateLimit)")
    public Object rateLimit(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        
        if (rateLimit == null) {
            return joinPoint.proceed();
        }

        String limitKey = getLimiterKey(method, rateLimit);
        RateLimiter rateLimiter = getRateLimiter(limitKey, rateLimit);

        // 尝试获取权限
        boolean permission = rateLimiter.acquirePermission();
        if (!permission) {
            log.warn("请求被限流 {}.{}", method.getDeclaringClass().getSimpleName(), method.getName());
            throw new BusinessException(ErrorCode.TOO_MANY_REQUESTS);
        }

        try {
            return joinPoint.proceed();
        } finally {
            // 可以在这里添加一些后置处理逻辑
        }
    }

    /**
     * 获取或创建限流器
     */
    private RateLimiter getRateLimiter(String key, RateLimit rateLimit) {
        return rateLimiterCache.computeIfAbsent(key, k -> {
            io.github.resilience4j.ratelimiter.RateLimiterConfig config = io.github.resilience4j.ratelimiter.RateLimiterConfig.custom()
                    .limitRefreshPeriod(Duration.ofSeconds(rateLimit.timeWindow()))
                    .limitForPeriod(rateLimit.limit())
                    .timeoutDuration(Duration.ofMillis(rateLimit.timeoutMs()))
                    .build();
            
            return rateLimiterRegistry.rateLimiter(key, config);
        });
    }

    /**
     * 生成限流器的唯一键
     */
    private String getLimiterKey(Method method, RateLimit rateLimit) {
        String key = rateLimit.key();
        if (key.isEmpty()) {
            key = method.getDeclaringClass().getSimpleName() + "." + method.getName();
        }
        return key;
    }
} 